#front matter
rm(list=ls())
library(spdep)
library(mgcv)
library(arm)
library(MCMCpack)
set.seed(314159265)
options(scipen=12)

#set working directory
#setwd('/Volumes/MONOGAN/stateAreal/mc/')

# Set-up areal structure using North Carolina example in spdep library 
example(nc.sids) 
n <- length(ncCR85_nb) 
coords <- coordinates(nc.sids) 
cards <- card(ncCR85_nb) 
col.w <- nb2listw(ncCR85_nb, style="W") #one weight matrix
col.w.alt <- nb2listw(ncCR85_nb, style="B") #another weight matrix
D<-array(unlist(lapply(col.w.alt$weights,sum))) #D matrix

#global parameters
ridge<- 0.97
constant<-1
beta1<-.7
beta2<-.5

#local parameter and number of treatments
psi<-seq(.1,.9,.1)
M<-length(psi)

#number of replicate samples
S<-150

#number of MCMC iterations
mcmc<-100000
burn<-10000

#output files
constant.car<-matrix(NA,nrow=S,ncol=M)
beta1.car<-matrix(NA,nrow=S,ncol=M)
beta2.car<-matrix(NA,nrow=S,ncol=M)
sigma.car<-matrix(NA,nrow=S,ncol=M)
delta.car<-matrix(NA,nrow=S,ncol=M)
psi.car<-matrix(NA,nrow=S,ncol=M)
real.psi<-matrix(NA,nrow=S,ncol=M)
total.sd<-matrix(NA,nrow=S,ncol=M)
total.var<-matrix(NA,nrow=S,ncol=M)
constant.lm<-matrix(NA,nrow=S,ncol=M)
beta1.lm<-matrix(NA,nrow=S,ncol=M)
beta2.lm<-matrix(NA,nrow=S,ncol=M)
sigma.lm<-matrix(NA,nrow=S,ncol=M)
phi.sd<-matrix(NA,nrow=S,ncol=M)
theta.sd<-matrix(NA,nrow=S,ncol=M)
constant.car.mae<-matrix(NA,nrow=S,ncol=M)
beta1.car.mae<-matrix(NA,nrow=S,ncol=M)
beta2.car.mae<-matrix(NA,nrow=S,ncol=M)
constant.lm.mae<-matrix(NA,nrow=S,ncol=M)
beta1.lm.mae<-matrix(NA,nrow=S,ncol=M)
beta2.lm.mae<-matrix(NA,nrow=S,ncol=M)

#function definitions
Lin.Model <- function(parm, X, y)
     {
     ### Parameters
     beta <- parm[1:dim(X)[2]]
     sigma <- exp(parm[dim(X)[2]+1])
     ### Log of Prior Densities
     sigma.prior <- log(dinvgamma(sigma, shape=1, scale=1))
     ### Log-Likelihood
     mu <- tcrossprod(X, t(beta))
     LL <- sum(log(dnorm(y, mean=mu, sd=sqrt(sigma))))
     ### Log-Posterior
    LP <- LL + sigma.prior
     return(LP)
     }

CAR.Model <- function(parm, X, y, W)
     {
     ### Parameters
     N<-dim(X)[1]
     beta <- parm[1:dim(X)[2]]
     sigma <- exp(parm[dim(X)[2]+1])
     delta <- exp(parm[dim(X)[2]+2])
     phi.vec<- parm[(dim(X)[2]+3):(dim(X)[2]+2+N)]
     Q<-length(phi.vec)

	 ### Log of Prior Densities
     sigma.prior <- log(dinvgamma(sigma, shape=1, scale=1))
     delta.prior <- log(dinvgamma(delta, shape=1, scale=1))

	### Log-Likelihood
     mu <- tcrossprod(X, t(beta))
     B<-diag(apply(W,1,sum))-ridge*W
     LL <- -(N/2)*log(sigma)-(Q/2)*log(delta)-((t(y-mu-phi.vec)%*%(y-mu-phi.vec))/(2*sigma))-((t(phi.vec)%*%B%*%phi.vec)/(2*delta)) 
          
     ### Log-Posterior
     LP <- LL + sigma.prior + delta.prior
     return(LP)
     }
     
#initial valuses
Initials.Lin<-c(.1,.1,.1,.1)
Initials.CAR<-c(.1,.1,.1,1,1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1,
.1,.1,.1,.1,.1,.1,.1,.1,.1,.1)


##########################  Model Specification  ##########################
#set up loop
M<-length(psi)
j<-1
#j<-9
while(j<=M){
	for(i in 1:S){
	#simulate variables
	x1<-rnorm(n,sd=1)
	x2.0<-rnorm(n,sd=1.5526751) #inflate initial variance ahead of inducing correlation
	tot<-psi[j]^2+(1-psi[j])^2
	phi.0 <- rnorm(n, sd=1.5526751*psi[j]/sqrt(tot)) #clustering standard deviation defined as a multiple of heterogeneous standard deviation (which itself is 1)
	theta <- rnorm(n, sd=(1-psi[j])/sqrt(tot))#does this work better?

	#create spatial error and dependent variable
	DrW <- diag(D) - ridge*listw2mat(col.w.alt) # I - rho * W
	CARcov<-solve(DrW)
	CARcovU = chol(CARcov) #  from  (D-rhoW)^ -1 
	phi<-t(CARcovU)%*%phi.0
	x2<-t(CARcovU)%*%x2.0
	
	y<-constant+beta1*x1+beta2*x2+phi+theta
	X<-cbind(1,x1,x2)

	#Fit Linear Model
	Lin.Fit <- MCMCmetrop1R(Lin.Model, theta.init=Initials.Lin, y=y, X=X, thin=1, mcmc=mcmc, burnin=burn, verbose=0,logfun=TRUE)
	constant.lm[i,j]<-summary(Lin.Fit)$quantiles[1,3]
	beta1.lm[i,j]<-summary(Lin.Fit)$quantiles[2,3]
	beta2.lm[i,j]<-summary(Lin.Fit)$quantiles[3,3]
	sigma.lm[i,j]<-exp(summary(Lin.Fit)$quantiles[4,3])

	#Fit CAR Model
	CAR.Fit <- MCMCmetrop1R(CAR.Model, theta.init=Initials.CAR, y=y, X=X, W=listw2mat(col.w), thin=1, mcmc=mcmc, burnin=burn, verbose=0,logfun=TRUE)
	constant.car[i,j]<-summary(CAR.Fit)$quantiles[1,3]
	beta1.car[i,j]<-summary(CAR.Fit)$quantiles[2,3]
	beta2.car[i,j]<-summary(CAR.Fit)$quantiles[3,3]
	beta.hat<- c(constant.car[i,j],beta1.car[i,j],beta2.car[i,j])
	phi.hat<-summary(CAR.Fit)$quantiles[6:105,3]
	theta.hat<- y-X%*%beta.hat -phi.hat
	sigma.car[i,j]<-sd(theta.hat)
	delta.car[i,j]<-sd(phi.hat)
	psi.car[i,j]<-(delta.car[i,j])/((delta.car[i,j])+(sigma.car[i,j]))
	real.psi[i,j]<-sd(phi)/(sd(phi)+sd(theta))
	phi.sd[i,j]<-sd(phi.0)
	theta.sd[i,j]<-sd(theta)
	total.sd[i,j]<-sd(phi+theta)
	total.var[i,j]<-var(phi+theta)
	
	#record absolute errors
	constant.car.mae[i,j]<-abs(constant.car[i,j]-constant)
	beta1.car.mae[i,j]<-abs(beta1.car[i,j]-beta1)
	beta2.car.mae[i,j]<-abs(beta2.car[i,j]-beta2)
	constant.lm.mae[i,j]<-abs(constant.lm[i,j]-constant)
	beta1.lm.mae[i,j]<-abs(beta1.lm[i,j]-beta1)
	beta2.lm.mae[i,j]<-abs(beta2.lm[i,j]-beta2)
	
	}
	
	#Cycle through outer loop
	cat("Round", j, "Complete \n")
	j<-j+1
}

#create summaries
constant.lm.sum<-apply(constant.lm,2,mean)
constant.car.sum<-apply(constant.car,2,mean)
beta1.lm.sum<-apply(beta1.lm,2,mean)
beta1.car.sum<-apply(beta1.car,2,mean)
beta2.lm.sum<-apply(beta2.lm,2,mean)
beta2.car.sum<-apply(beta2.car,2,mean)
psi.car.sum<-apply(psi.car,2,mean)
real.psi.sum<-apply(real.psi,2,mean)
total.sd.sum<-apply(total.sd,2,mean)
total.var.sum<-apply(total.var,2,mean)
phi.sd.sum<-apply(phi.sd,2,mean)
theta.sd.sum<-apply(theta.sd,2,mean)
mae.constant.lm.sum<-apply(constant.lm.mae,2,mean)
mae.constant.car.sum<-apply(constant.car.mae,2,mean)
mae.beta1.lm.sum<-apply(beta1.lm.mae,2,mean)
mae.beta1.car.sum<-apply(beta1.car.mae,2,mean)
mae.beta2.lm.sum<-apply(beta2.lm.mae,2,mean)
mae.beta2.car.sum<-apply(beta2.car.mae,2,mean)



#Graph the results 
#Draw Figure 3.a
#pdf("correlOther.pdf")
pdf("figure3a.pdf")
par(mar=c(5.1,4.5,4.1,1.9))
plot(x=psi,y=mae.beta1.lm.sum,type='h',xlab=expression(psi),ylab=expression(paste("Mean Absolute Error for ", beta[1])),col='blue',lty=2,lwd=2,xlim=c(0,1),ylim=c(0,.15),axes=F)
axis(1,at=seq(0,1,.1));axis(2);box()
abline(h=0,col='gray60')
par(new=T)
plot(x=psi+.01, y=mae.beta1.car.sum, type='h',xlab="",ylab="",axes=F, col='red', lty=1,lwd=2,xlim=c(0,1),ylim=c(0,.15))
legend(x=0,y=.15,legend=c("Simple Linear","CAR"), lty=c(2,1),lwd=c(2,2),col=c('blue','red'))
dev.off()

#Draw Figure 3.b
#pdf("correlCorrel.pdf")
pdf("figure3b.pdf")
par(mar=c(5.1,4.5,4.1,1.9))
plot(x=psi,y=mae.beta2.lm.sum,type='h',xlab=expression(psi),ylab=expression(paste("Mean Absolute Error for ", beta[2])),col='blue',lty=2,lwd=2,xlim=c(0,1),ylim=c(0,.2),axes=F)
axis(1,at=seq(0,1,.1));axis(2);box()
abline(h=0,col='gray60')
par(new=T)
plot(x=psi+.01, y=mae.beta2.car.sum, type='h',xlab="",ylab="",axes=F, col='red', lty=1,lwd=2,xlim=c(0,1),ylim=c(0,.2))
legend(x=0,y=.2,legend=c("Simple Linear","CAR"), lty=c(2,1),lwd=c(2,2),col=c('blue','red'))
dev.off()

